import os
import json
import pickle
import numpy as np
from pymilvus import (
    connections,
    utility,
    FieldSchema, CollectionSchema, DataType,
    Collection,
)
from typing import List, Dict, Any, Union, Tuple
from datetime import datetime
import shutil
from tqdm import tqdm
import glob
from PIL import Image

from models.clip_model import ClipModel

class KnowledgeBaseLite:
    """知识库管理类，使用Milvus Lite进行本地向量存储"""
    
    def __init__(self, base_path: str = "./knowledge_base", clip_model: ClipModel = None):
        """
        初始化知识库
        
        Args:
            base_path: 知识库存储路径
            clip_model: CLIP模型实例
        """
        self.base_path = base_path
        self.clip_model = clip_model
        
        # 创建知识库目录
        os.makedirs(base_path, exist_ok=True)
        
        # 知识库元数据
        self.kb_info = {
            "product_types": [],
            "created_at": None,
            "updated_at": None,
            "total_images": 0,
            "features_dim": 0,
            "clip_model": None
        }
        
        # 初始化
        self.collection = None
        self.image_paths = []
        self.image_labels = []
        
        # 连接到Milvus Lite
        try:
            connections.connect(alias="default", uri="lite://.milvus_lite")
            print("已连接到Milvus Lite本地向量库")
        except Exception as e:
            print(f"连接Milvus Lite失败: {str(e)}")
        
    def load_or_create(self, product_name: str) -> bool:
        """
        加载或创建产品知识库
        
        Args:
            product_name: 产品名称
            
        Returns:
            是否成功加载或创建
        """
        product_path = os.path.join(self.base_path, product_name)
        os.makedirs(product_path, exist_ok=True)
        
        # 检查是否存在元数据文件
        metadata_path = os.path.join(product_path, "metadata.json")
        paths_path = os.path.join(product_path, "image_paths.pkl")
        labels_path = os.path.join(product_path, "image_labels.pkl")
        
        collection_name = f"defect_kb_{product_name}"
        
        # 如果存在元数据文件，则加载
        if os.path.exists(metadata_path) and utility.has_collection(collection_name):
            try:
                # 加载元数据
                with open(metadata_path, 'r', encoding='utf-8') as f:
                    self.kb_info = json.load(f)
                
                # 加载图像路径和标签
                with open(paths_path, 'rb') as f:
                    self.image_paths = pickle.load(f)
                
                with open(labels_path, 'rb') as f:
                    self.image_labels = pickle.load(f)
                
                # 加载Milvus集合
                self.collection = Collection(collection_name)
                self.collection.load()
                
                print(f"已加载知识库: {product_name}")
                print(f"包含图片数: {self.kb_info['total_images']}")
                return True
            
            except Exception as e:
                print(f"加载知识库失败: {str(e)}")
                return False
        
        return True
    
    def build_from_folder(self, product_name: str, data_dir: str, batch_size: int = 32) -> bool:
        """
        从文件夹构建知识库
        
        Args:
            product_name: 产品名称
            data_dir: 数据文件夹路径（应包含良品和缺陷子文件夹）
            batch_size: 批处理大小
            
        Returns:
            是否成功构建
        """
        if self.clip_model is None:
            print("错误: 请先初始化CLIP模型")
            return False
        
        # 创建产品知识库目录
        product_path = os.path.join(self.base_path, product_name)
        os.makedirs(product_path, exist_ok=True)
        
        # 获取所有子文件夹（代表不同缺陷类型）
        subdirs = [d for d in os.listdir(data_dir) if os.path.isdir(os.path.join(data_dir, d))]
        
        if not subdirs:
            print(f"警告: {data_dir} 中没有找到子文件夹")
            return False
        
        all_features = []
        self.image_paths = []
        self.image_labels = []
        
        total_images = 0
        for subdir in subdirs:
            subdir_path = os.path.join(data_dir, subdir)
            image_files = []
            
            # 获取所有支持的图片文件
            for ext in ['png', 'jpg', 'jpeg']:
                image_files.extend(glob.glob(os.path.join(subdir_path, f"*.{ext}")))
            
            if not image_files:
                print(f"警告: {subdir_path} 中没有找到图片")
                continue
            
            print(f"处理 {subdir} 中的 {len(image_files)} 张图片...")
            
            # 批处理提取特征
            for i in range(0, len(image_files), batch_size):
                batch = image_files[i:i+batch_size]
                features, valid_paths = self.clip_model.extract_features(batch)
                
                if len(valid_paths) > 0:
                    all_features.append(features)
                    self.image_paths.extend(valid_paths)
                    self.image_labels.extend([subdir] * len(valid_paths))
                    total_images += len(valid_paths)
        
        if total_images == 0:
            print("错误: 没有成功处理任何图片")
            return False
        
        # 合并所有特征
        all_features = np.vstack(all_features)
        features_dim = all_features.shape[1]
        
        # 创建Milvus集合
        collection_name = f"defect_kb_{product_name}"
        
        # 如果集合已存在，则删除
        if utility.has_collection(collection_name):
            utility.drop_collection(collection_name)
        
        # 定义集合架构
        fields = [
            FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
            FieldSchema(name="image_path", dtype=DataType.VARCHAR, max_length=512),
            FieldSchema(name="label", dtype=DataType.VARCHAR, max_length=64),
            FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=features_dim)
        ]
        schema = CollectionSchema(fields, "产品表面缺陷知识库")
        self.collection = Collection(collection_name, schema)
        
        # 创建索引
        index_params = {
            "metric_type": "IP",  # 内积，对应余弦相似度
            "index_type": "IVF_FLAT",
            "params": {"nlist": 128}  # Lite模式下使用较小的nlist
        }
        self.collection.create_index(field_name="embedding", index_params=index_params)
        
        # 插入数据
        data = [
            self.image_paths,
            self.image_labels,
            all_features.tolist()
        ]
        
        self.collection.insert(data)
        self.collection.load()
        
        # 更新知识库信息
        self.kb_info = {
            "product_name": product_name,
            "created_at": datetime.now().isoformat(),
            "updated_at": datetime.now().isoformat(),
            "total_images": total_images,
            "features_dim": features_dim,
            "clip_model": self.clip_model.model_name,
            "defect_types": list(set(self.image_labels))
        }
        
        # 保存索引和元数据
        self._save_knowledge_base(product_name)
        
        print(f"知识库构建完成: {product_name}")
        print(f"包含图片数: {total_images}")
        print(f"缺陷类型: {list(set(self.image_labels))}")
        
        return True
    
    def add_images(self, product_name: str, image_paths: List[str], labels: List[str]) -> bool:
        """
        向现有知识库中添加图片
        
        Args:
            product_name: 产品名称
            image_paths: 图片路径列表
            labels: 对应的标签列表
            
        Returns:
            是否成功添加
        """
        if len(image_paths) != len(labels):
            print("错误: 图片路径和标签数量不匹配")
            return False
        
        if self.clip_model is None:
            print("错误: 请先初始化CLIP模型")
            return False
        
        # 加载现有知识库
        if not self.load_or_create(product_name):
            print(f"警告: 无法加载知识库 {product_name}，将创建新知识库")
        
        # 提取特征
        features, valid_paths = self.clip_model.extract_features(image_paths)
        
        if len(valid_paths) == 0:
            print("错误: 没有成功处理任何图片")
            return False
        
        collection_name = f"defect_kb_{product_name}"
        
        # 如果集合不存在，创建新集合
        if not utility.has_collection(collection_name):
            features_dim = features.shape[1]
            
            # 定义集合架构
            fields = [
                FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
                FieldSchema(name="image_path", dtype=DataType.VARCHAR, max_length=512),
                FieldSchema(name="label", dtype=DataType.VARCHAR, max_length=64),
                FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=features_dim)
            ]
            schema = CollectionSchema(fields, "产品表面缺陷知识库")
            self.collection = Collection(collection_name, schema)
            
            # 创建索引
            index_params = {
                "metric_type": "IP",  # 内积，对应余弦相似度
                "index_type": "IVF_FLAT",
                "params": {"nlist": 128}  # Lite模式下使用较小的nlist
            }
            self.collection.create_index(field_name="embedding", index_params=index_params)
            
            self.kb_info["features_dim"] = features_dim
            self.kb_info["clip_model"] = self.clip_model.model_name
        
        # 插入新数据
        valid_labels = [labels[image_paths.index(path)] for path in valid_paths]
        
        data = [
            valid_paths,
            valid_labels,
            features.tolist()
        ]
        
        self.collection.insert(data)
        
        # 更新图片路径和标签
        self.image_paths.extend(valid_paths)
        self.image_labels.extend(valid_labels)
        
        # 更新知识库信息
        self.kb_info["total_images"] = len(self.image_paths)
        self.kb_info["updated_at"] = datetime.now().isoformat()
        if "defect_types" in self.kb_info:
            defect_types = set(self.kb_info["defect_types"])
            defect_types.update(valid_labels)
            self.kb_info["defect_types"] = list(defect_types)
        else:
            self.kb_info["defect_types"] = list(set(valid_labels))
        
        # 保存索引和元数据
        self._save_knowledge_base(product_name)
        
        print(f"已向知识库 {product_name} 添加 {len(valid_paths)} 张图片")
        return True
    
    def search(self, query_image_path: str, top_k: int = 5) -> List[Dict]:
        """
        搜索最相似的图片
        
        Args:
            query_image_path: 查询图片路径
            top_k: 返回结果数量
            
        Returns:
            最相似图片的信息列表
        """
        if self.collection is None:
            print("错误: 知识库未加载")
            return []
        
        if self.clip_model is None:
            print("错误: CLIP模型未初始化")
            return []
        
        # 提取查询图片特征
        query_features, _ = self.clip_model.extract_features([query_image_path])
        
        if query_features.shape[0] == 0:
            print(f"错误: 无法处理查询图片 {query_image_path}")
            return []
        
        # 搜索相似图片
        search_params = {
            "metric_type": "IP",
            "params": {"nprobe": 16}
        }
        
        results = self.collection.search(
            data=[query_features.tolist()[0]],
            anns_field="embedding",
            param=search_params,
            limit=top_k,
            output_fields=["image_path", "label"]
        )
        
        # 整理结果
        result_list = []
        for hits in results:
            for hit in hits:
                result_list.append({
                    "similarity": float(hit.score),
                    "path": hit.entity.get("image_path"),
                    "label": hit.entity.get("label")
                })
        
        return result_list
    
    def get_defect_types(self, product_name: str) -> List[str]:
        """
        获取产品的缺陷类型列表
        
        Args:
            product_name: 产品名称
            
        Returns:
            缺陷类型列表
        """
        if not self.load_or_create(product_name):
            return []
        
        if "defect_types" in self.kb_info:
            return self.kb_info["defect_types"]
        else:
            return list(set(self.image_labels))
    
    def get_product_list(self) -> List[str]:
        """
        获取所有产品列表
        
        Returns:
            产品名称列表
        """
        products = []
        if os.path.exists(self.base_path):
            products = [d for d in os.listdir(self.base_path) 
                       if os.path.isdir(os.path.join(self.base_path, d))]
        
        return products
    
    def delete_knowledge_base(self, product_name: str) -> bool:
        """
        删除知识库
        
        Args:
            product_name: 产品名称
            
        Returns:
            是否成功删除
        """
        product_path = os.path.join(self.base_path, product_name)
        collection_name = f"defect_kb_{product_name}"
        
        try:
            # 删除Milvus集合
            if utility.has_collection(collection_name):
                utility.drop_collection(collection_name)
            
            # 删除本地文件
            if os.path.exists(product_path):
                shutil.rmtree(product_path)
                
            print(f"已删除知识库: {product_name}")
            
            # 重置当前加载的知识库（如果是被删除的知识库）
            if self.kb_info.get("product_name") == product_name:
                self.collection = None
                self.image_paths = []
                self.image_labels = []
                self.kb_info = {"product_types": []}
            
            return True
        except Exception as e:
            print(f"删除知识库失败: {str(e)}")
            return False
    
    def _save_knowledge_base(self, product_name: str) -> None:
        """
        保存知识库相关文件
        
        Args:
            product_name: 产品名称
        """
        product_path = os.path.join(self.base_path, product_name)
        os.makedirs(product_path, exist_ok=True)
        
        # 保存元数据
        metadata_path = os.path.join(product_path, "metadata.json")
        with open(metadata_path, 'w', encoding='utf-8') as f:
            json.dump(self.kb_info, f, ensure_ascii=False, indent=2)
        
        # 保存图像路径和标签
        paths_path = os.path.join(product_path, "image_paths.pkl")
        with open(paths_path, 'wb') as f:
            pickle.dump(self.image_paths, f)
        
        labels_path = os.path.join(product_path, "image_labels.pkl")
        with open(labels_path, 'wb') as f:
            pickle.dump(self.image_labels, f) 